package org.goobs.applications;

import java.util.ArrayList;
import java.util.HashSet;

import org.goobs.csp.*;

public class QueensCSP {

	private static int numQueens = 10;
	
	private static class PosUnique extends UniquenessConstraint{
		public PosUnique(CSPVariable var) {
			super(var);
		}
		@Override
		public boolean otherVarUnique(CSPVariable otherVar) {
			Location otherLoc = (Location) otherVar.getValue();
			Location thisLoc = (Location) this.var.getValue();
			return thisLoc.x != otherLoc.x && thisLoc.y != otherLoc.y;
		}
		@Override
		public boolean thisVarQualifies() {
			return true;
		}
	}
	
	private static class DiagUnique extends UniquenessConstraint{
		public DiagUnique(CSPVariable var) {
			super(var);
		}
		@Override
		public boolean otherVarUnique(CSPVariable otherVar) {
			Location otherLoc = (Location) otherVar.getValue();
			Location thisLoc = (Location) this.var.getValue();
			int rise = thisLoc.y - otherLoc.y;
			int run = thisLoc.x - otherLoc.x;
			if(Math.abs(rise) == Math.abs(run)){
				return false;
			}else{
				return true;
			}
		}
		@Override
		public boolean thisVarQualifies() {
			return true;
		}
		
	}
	
	private static class LocConstraint extends BinaryConstraint <Integer, Integer>{
		public LocConstraint(CSPVariable <Integer> a, CSPVariable <Integer> b) {
			super(a, b, true);
		}
		@Override
		public boolean isSatisfied(CSPVariable tail, CSPVariable head) {
			int a = Math.abs( (Integer) tail.getData() );
			int b = Math.abs( (Integer) head.getData() );
			int aAssign = (Integer) tail.getValue();
			int bAssign = (Integer) head.getValue();
			
			if(aAssign == b){
				return bAssign == a;
			}else{
				return true;
			}
		}
	}
	
	private static class DiagUniqueRowCol extends BinaryConstraint{	
		public DiagUniqueRowCol(CSPVariable a, CSPVariable b) {
			super(a, b, false);
		}
		@Override
		public boolean isSatisfied(CSPVariable tail, CSPVariable head) {
			int a = Math.abs( (Integer) tail.getData() );
			int b = Math.abs( (Integer) head.getData() );
			int aAssign = (Integer) tail.getValue();
			int bAssign = (Integer) head.getValue();
			
			int rise = bAssign - aAssign;
			int run = b - a;
			if( Math.abs(rise) == Math.abs(run) ){
				return false;
			}else{
//				System.out.println(tail + "\n" + head + "\n");
				return true;
			}
		}
		
		
	}
	
	private static class Location{
		private int x;
		private int y;
		public String toString(){
			return "(" + x + "," + y + ")";
		}
	}
	
	
	private static void queensVar(){
		System.out.println("Setting up CSP...");
		CSP queens = new CSP();
		
		System.out.println("   creating domain...");
		final HashSet <Location> locSet = new HashSet <Location> ();
		for(int i=0; i<numQueens; i++){
			for(int j=0; j<numQueens; j++){
				Location loc = new Location();
				loc.x = i;
				loc.y = j;
				locSet.add(loc);
			}
		}
		
		System.out.println("   adding variables and basic constraints...");
		for(int i=0; i<numQueens; i++){
			System.out.println("      Queen " + i);
			CSPVariable <Location> var = new CSPVariable <Location> (locSet, "Queen " + i);
			queens.addVariable(var);
			queens.addConstraint(new PosUnique(var));
			queens.addConstraint(new DiagUnique(var));
		}
		
		System.out.println();
		queens.printState();
		
		System.out.println("Running CSP...");
		if(!queens.register()){
			System.out.println("   COULD NOT REGISTER CSP");
			System.exit(1);
		}
		if(!queens.findSolution()){
			System.out.println("NO SOLUTION");
		}else{
			queens.printState();
			int i=1;
			while(queens.findSolution()){
				i++;
			}
			System.out.println("total solution count: " + i);
		}
	}
	
	
	private static void rowColVar(){
		System.out.println("Setting up CSP...");
		CSP queens = new CSP();
		
		System.out.println("   creating domain...");
		final HashSet <Integer> posSet = new HashSet <Integer> ();
		for(int i=0; i<numQueens; i++){
			posSet.add(i+1);
		}
		
		System.out.println("   adding variables and basic constraints...");
		ArrayList <CSPVariable <Integer>> rows = new ArrayList <CSPVariable <Integer>> ();
		ArrayList <CSPVariable <Integer>> cols = new ArrayList <CSPVariable <Integer>> ();
		for(int i=0; i<numQueens; i++){
			System.out.println("      Row" + i);
			CSPVariable <Integer> var = new CSPVariable <Integer> (posSet, i+1);
			queens.addVariable(var);
			rows.add(var);
			System.out.println("      Col" + i);
			var = new CSPVariable <Integer> (posSet, -(i+1));
			queens.addVariable(var);
			cols.add(var);
		}
		
		System.out.println("   adding compex constraints...");
		for(CSPVariable <Integer> row : rows){
			for(CSPVariable <Integer> col : cols){
				LocConstraint cons = new LocConstraint(row, col);
				queens.addConstraint(cons);
			}
			for(CSPVariable <Integer> row2 : rows){
				if(row != row2){
					DiagUniqueRowCol cons = new DiagUniqueRowCol(row, row2);
					queens.addConstraint(cons);
				}
			}
		}
		
		
		
		System.out.println();
		queens.printState();
		
		System.out.println("Running CSP...");
		if(!queens.register()){
			System.out.println("   COULD NOT REGISTER CSP");
			System.exit(1);
		}
		if(!queens.findSolution()){
			System.out.println("NO SOLUTION");
		}else{
			queens.printState();
//			int i=1;
//			while(queens.findSolution()){
//				i++;
//			}
//			System.out.println("total solution count: " + i);
		}
	}
	
	public static void main(String[] args){
		rowColVar();
	}

}
